{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 无迹卡尔曼滤波\n",
    "使用UKF进行移动机器人的位姿预测，见相关[博客](https://blog.csdn.net/weixin_42301220/article/details/124708187?csdn_share_tail=%7B%22type%22%3A%22blog%22%2C%22rType%22%3A%22article%22%2C%22rId%22%3A%22124708187%22%2C%22source%22%3A%22weixin_42301220%22%7D&ctrtid=B4W0n)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "import PyQt5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "# %matplotlib inline\n",
    "%matplotlib qt5\n",
    "\n",
    "# # set up matplotlib\n",
    "# is_ipython = 'inline' in matplotlib.get_backend()\n",
    "# if is_ipython:\n",
    "#     from IPython import display\n",
    "\n",
    "plt.ion()\n",
    "\n",
    "import scipy\n",
    "import numpy as np\n",
    "from scipy.spatial.transform import Rotation as Rot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "设置过程噪声和观测噪声的协方差矩阵"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "Q = np.matrix([[0.1, 0, 0, 0],  # variance of location on x-axis\n",
    "        [0, 0.1, 0, 0],  # variance of location on y-axis\n",
    "        [0, 0, np.deg2rad(1.0), 0],  # variance of yaw angle\n",
    "        [0, 0, 0, 1.0]])**2  # variance of velocity\n",
    "\n",
    "R = np.matrix([[1.0,0.0],[0.0,1.0]])**2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "仿真参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "#  Simulation parameter\n",
    "INPUT_NOISE = np.diag([1.0, np.deg2rad(30.0)]) ** 2\n",
    "GPS_NOISE = np.diag([0.5, 0.5]) ** 2\n",
    "\n",
    "DT = 0.1  # time tick [s]\n",
    "SIM_TIME = 50.0  # simulation time [s]\n",
    "\n",
    "show_animation = True\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "UKF参数\n",
    "\n",
    "使用比例无迹变换\n",
    "\n",
    "![在这里插入图片描述](https://img-blog.csdnimg.cn/7df0f05a8c624ef497ccca3e89711562.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "#  UKF Parameter\n",
    "ALPHA = 0.001\n",
    "BETA = 2\n",
    "KAPPA = 0\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "UKF初始化：参数选择与权重计算\n",
    "![在这里插入图片描述](https://img-blog.csdnimg.cn/ad4e2a00248b484e91804d018b517e8c.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup_ukf(nx):\n",
    "    lamb = ALPHA ** 2 * (nx + KAPPA) - nx\n",
    "    # calculate weights\n",
    "    wm = [lamb / (lamb + nx)]\n",
    "    wc = [(lamb / (lamb + nx)) + (1 - ALPHA ** 2 + BETA)]\n",
    "    for i in range(2 * nx):\n",
    "        wm.append(1.0 / (2 * (nx + lamb)))\n",
    "        wc.append(1.0 / (2 * (nx + lamb)))\n",
    "    gamma = math.sqrt(nx + lamb)\n",
    "\n",
    "    wm = np.array([wm])\n",
    "    wc = np.array([wc])\n",
    "\n",
    "    return wm, wc, gamma\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "状态向量\n",
    "\n",
    "x_t = [x,y,yaw,v]\n",
    "\n",
    "控制输入为\n",
    "\n",
    "u = [v,w]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "控制输入表示"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "def control_input():\n",
    "    v = 1.0  # 线速度，m/s\n",
    "    yaw_rate = 0.1  #角速度，rad/s\n",
    "    u = np.matrix([\n",
    "            [v],\n",
    "            [yaw_rate]\n",
    "            ])\n",
    "    return u"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "观测模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "def observation_model(x):\n",
    "    H = np.matrix([[1.0,0,0,0],[0,1.0,0,0]])\n",
    "    z = H@x\n",
    "    return z\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "过程模型（状态方程）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "def motion_model(x,u):\n",
    "    F = np.matrix([\n",
    "            [1.0,0,0,0],\n",
    "            [0,1.0,0,0],\n",
    "            [0,0,1.0,0],\n",
    "            [0,0,0,0]\n",
    "            ])\n",
    "    B = np.matrix([\n",
    "            [DT*np.cos(x[2,0]),0],\n",
    "            [DT*np.sin(x[2,0]),0],\n",
    "            [0,DT],\n",
    "            [1.0,0]\n",
    "            ])\n",
    "    x = F@x+B@u\n",
    "    return x\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "观测向量获取"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "def observation(x_true, x_d,u):\n",
    "    x_true = motion_model(x_true,u)\n",
    "    # add noise to gps x-y\n",
    "    z = observation_model(x_true) + GPS_NOISE @ np.random.randn(2, 1)\n",
    "\n",
    "    # add noise to input\n",
    "    ud = u + INPUT_NOISE @ np.random.randn(2, 1)\n",
    "\n",
    "    xd = motion_model(x_d, ud)\n",
    "\n",
    "    return x_true, z, xd, ud\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "sigma点采样\n",
    "\n",
    "![在这里插入图片描述](https://img-blog.csdnimg.cn/6c6d31779f674b8aa27c0d2b5a351065.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_sigma_points(xEst, PEst, gamma):\n",
    "    sigma = xEst\n",
    "    Psqrt = scipy.linalg.sqrtm(PEst)\n",
    "    n = len(xEst[:, 0])\n",
    "    # Positive direction\n",
    "    for i in range(n):\n",
    "        sigma = np.hstack((sigma, xEst + gamma * Psqrt[:, i:i + 1]))\n",
    "\n",
    "    # Negative direction\n",
    "    for i in range(n):\n",
    "        sigma = np.hstack((sigma, xEst - gamma * Psqrt[:, i:i + 1]))\n",
    "\n",
    "    return sigma\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "状态转移非线性变换\n",
    "\n",
    "![在这里插入图片描述](https://img-blog.csdnimg.cn/28c6485d04ee432f849ed12d95b8f455.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_sigma_motion(sigma, u):\n",
    "    \"\"\"\n",
    "        Sigma Points prediction with motion model\n",
    "    \"\"\"\n",
    "    for i in range(sigma.shape[1]):\n",
    "        sigma[:, i:i + 1] = motion_model(sigma[:, i:i + 1], u)\n",
    "\n",
    "    return sigma"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "观测非线性变换\n",
    "![在这里插入图片描述](https://img-blog.csdnimg.cn/3e97c759dadf4915a8725e26723e91ac.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_sigma_observation(sigma):\n",
    "    \"\"\"\n",
    "        Sigma Points prediction with observation model\n",
    "    \"\"\"\n",
    "    for i in range(sigma.shape[1]):\n",
    "        sigma[0:2, i] = observation_model(sigma[:, i])\n",
    "\n",
    "    sigma = sigma[0:2, :]\n",
    "\n",
    "    return sigma\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " 加权计算 k 时刻状态量的先验概率分布\n",
    " ![在这里插入图片描述](https://img-blog.csdnimg.cn/5112c5661d54476ab1ca93b585ff8c30.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_sigma_covariance(x, sigma, wc, Pi):\n",
    "    nSigma = sigma.shape[1]\n",
    "    d = sigma - x[0:sigma.shape[0]]\n",
    "    P = Pi\n",
    "    for i in range(nSigma):\n",
    "        P = P + wc[0, i] * d[:, i:i + 1] @ d[:, i:i + 1].T\n",
    "    return P\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "计算状态量与观测量的互协方差\n",
    "\n",
    "![在这里插入图片描述](https://img-blog.csdnimg.cn/a9aa0a94ca69404ba47cf2c90066ec4a.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_pxz(sigma, x, z_sigma, zb, wc):\n",
    "    nSigma = sigma.shape[1]\n",
    "    dx = sigma - x\n",
    "    dz = z_sigma - zb[0:2]\n",
    "    P = np.zeros((dx.shape[0], dz.shape[0]))\n",
    "\n",
    "    for i in range(nSigma):\n",
    "        P = P + wc[0, i] * dx[:, i:i + 1] @ dz[:, i:i + 1].T\n",
    "\n",
    "    return P\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "UKF估计"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "代码符号说明\n",
    "- 第1个sigma：![在这里插入图片描述](https://img-blog.csdnimg.cn/de69dc49635c4139bbda6883588d72ba.png)\n",
    "\n",
    "- 第2个sigma：![在这里插入图片描述](https://img-blog.csdnimg.cn/7bdb084770ee43758f98654b61895bcc.png)\n",
    "\n",
    "- xPred：![在这里插入图片描述](https://img-blog.csdnimg.cn/330e6716828548c78cc39bbb7f5803fc.png)\n",
    "\n",
    "- PPred：![在这里插入图片描述](https://img-blog.csdnimg.cn/cf3a523bb046484abbe9d110e9c7b3ca.png)\n",
    "\n",
    "- 第3个sigma：![在这里插入图片描述](https://img-blog.csdnimg.cn/8cd1bcd7ea5449f5b62ccd85a3f29917.png)\n",
    "\n",
    "- z_sigma：![在这里插入图片描述](https://img-blog.csdnimg.cn/3af2cd6f1fb74c84a064899fc760fa05.png)\n",
    "\n",
    "- zb：![在这里插入图片描述](https://img-blog.csdnimg.cn/d9a560bb9e58448598f9f6f431abbe34.png)\n",
    "\n",
    "- st：![在这里插入图片描述](https://img-blog.csdnimg.cn/8c8d1b74d1c94007b778cec6bf2d4cef.png)\n",
    "\n",
    "- Pxz：![在这里插入图片描述](https://img-blog.csdnimg.cn/e22ba23b218b48f8aa2bfe7ca4d92e22.png)\n",
    "\n",
    "- K：卡尔曼增益![在这里插入图片描述](https://img-blog.csdnimg.cn/5315544217454febac7535f777823ee5.png)\n",
    "\n",
    "- z：![在这里插入图片描述](https://img-blog.csdnimg.cn/ebe9e42b0487460fa36f031bca28077a.png)\n",
    "\n",
    "- y：![在这里插入图片描述](https://img-blog.csdnimg.cn/053c058490874d4fae78be4eb82cbfdc.png)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ukf_estimation(xEst, PEst, z, u, wm, wc, gamma):\n",
    "    #  Predict\n",
    "    \n",
    "    ## 1.sigma采样\n",
    "    sigma = generate_sigma_points(xEst, PEst, gamma)\n",
    "    \n",
    "    ## 2.状态转移非线性变换\n",
    "    sigma = predict_sigma_motion(sigma, u)\n",
    "    \n",
    "    ## 3.加权计算 k 时刻状态量的先验概率分布\n",
    "    xPred = (wm @ sigma.T).T\n",
    "    PPred = calc_sigma_covariance(xPred, sigma, wc, Q)\n",
    "\n",
    "    #  Update\n",
    "\n",
    "    ## 4.对 k 时刻状态量的先验概率分布进行 sigma 采样\n",
    "    sigma = generate_sigma_points(xPred, PPred, gamma)\n",
    "    \n",
    "\n",
    "    \n",
    "    ## 5.观测非线性变换\n",
    "    z_sigma = predict_sigma_observation(sigma)\n",
    "    \n",
    "    ## 6.加权计算k时刻观测量 Z_{k}的概率分布\n",
    "    # zb = (wm @ sigma.T).T\n",
    "    zb = (wm @ z_sigma.T).T\n",
    "    st = calc_sigma_covariance(zb, z_sigma, wc, R)\n",
    "    \n",
    "    ## 7. 计算状态量与观测量的互协方差\n",
    "    Pxz = calc_pxz(sigma, xPred, z_sigma, zb, wc)\n",
    "    \n",
    "    ## 8. 计算卡尔曼增益\n",
    "    K = Pxz @ np.linalg.inv(st)\n",
    "    \n",
    "    ## 计算 k 时刻状态量的后验概率分布\n",
    "    zPred = observation_model(xPred)\n",
    "    # y = z - zPred\n",
    "    y = z - zb\n",
    "    xEst = xPred + K @ y\n",
    "    PEst = PPred - K @ st @ K.T\n",
    "\n",
    "    return xEst, PEst\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "画图"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_covariance_ellipse(xEst, PEst):  # pragma: no cover\n",
    "    \"\"\"The blue line is true trajectory, the black line is dead reckoning trajectory,\n",
    "    the green point is positioning observation (ex. GPS), and the red line is estimated trajectory with EKF.\n",
    "\n",
    "    The red ellipse is estimated covariance ellipse with UKF.\n",
    "\n",
    "    Args:\n",
    "        xEst (_type_): _description_\n",
    "        PEst (_type_): _description_\n",
    "    \"\"\"\n",
    "    Pxy = PEst[0:2, 0:2]\n",
    "    eigval, eigvec = np.linalg.eig(Pxy)\n",
    "\n",
    "    if eigval[0] >= eigval[1]:\n",
    "        bigind = 0\n",
    "        smallind = 1\n",
    "    else:\n",
    "        bigind = 1\n",
    "        smallind = 0\n",
    "\n",
    "    t = np.arange(0, 2 * math.pi + 0.1, 0.1)\n",
    "    a = math.sqrt(eigval[bigind])\n",
    "    b = math.sqrt(eigval[smallind])\n",
    "    x = [a * math.cos(it) for it in t]\n",
    "    y = [b * math.sin(it) for it in t]\n",
    "    angle = math.atan2(eigvec[1, bigind], eigvec[0, bigind])\n",
    "    rot = Rot.from_euler('z', angle).as_matrix()[0:2, 0:2]\n",
    "    fx = rot @ np.array([x, y])\n",
    "    px = np.array(fx[0, :] + xEst[0, 0]).flatten()\n",
    "    py = np.array(fx[1, :] + xEst[1, 0]).flatten()\n",
    "    plt.plot(px, py, \"--r\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "主函数入口"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    nx = 4  # State Vector [x y yaw v]'\n",
    "    xEst = np.zeros((nx, 1))\n",
    "    xTrue = np.zeros((nx, 1))\n",
    "    PEst = np.eye(nx)\n",
    "    xDR = np.zeros((nx, 1))  # Dead reckoning\n",
    "\n",
    "    wm, wc, gamma = setup_ukf(nx)\n",
    "\n",
    "    # history\n",
    "    hxEst = xEst\n",
    "    hxTrue = xTrue\n",
    "    hxDR = xTrue\n",
    "    hz = np.zeros((2, 1))\n",
    "\n",
    "    time = 0.0\n",
    "\n",
    "    while SIM_TIME >= time:\n",
    "        time += DT\n",
    "        u = control_input()\n",
    "\n",
    "        xTrue, z, xDR, ud = observation(xTrue, xDR, u)\n",
    "\n",
    "        xEst, PEst = ukf_estimation(xEst, PEst, z, ud, wm, wc, gamma)\n",
    "\n",
    "        # store data history\n",
    "        hxEst = np.hstack((hxEst, xEst))\n",
    "        hxDR = np.hstack((hxDR, xDR))\n",
    "        hxTrue = np.hstack((hxTrue, xTrue))\n",
    "        hz = np.hstack((hz, z))\n",
    "\n",
    "        if show_animation:\n",
    "                plt.cla()\n",
    "                # for stopping simulation with the esc key.\n",
    "                plt.gcf().canvas.mpl_connect('key_release_event',lambda event: [exit(0) if event.key == 'escape' else None])\n",
    "                plt.plot(hz[0, :], hz[1, :], \".g\")\n",
    "                plt.plot(np.array(hxTrue[0, :]).flatten(),\n",
    "                        np.array(hxTrue[1, :]).flatten(), \"-b\")\n",
    "                plt.plot(np.array(hxDR[0, :]).flatten(),\n",
    "                        np.array(hxDR[1, :]).flatten(), \"-k\")\n",
    "                plt.plot(np.array(hxEst[0, :]).flatten(),\n",
    "                        np.array(hxEst[1, :]).flatten(), \"-r\")\n",
    "                plot_covariance_ellipse(xEst, PEst)\n",
    "                plt.axis(\"equal\")\n",
    "                plt.grid(True)\n",
    "                plt.pause(0.001)\n",
    "        plt.savefig('./result.png')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "main()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "0c7484b3574347463e16b31029466871583b0d4e5c4ad861e8848f2d3746b4de"
  },
  "kernelspec": {
   "display_name": "Python 3.8.12 ('gobigger')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
